from spacemit_cv import AGVDetection
import cv2
import argparse
import numpy as np

#from tracker.byte_tracker import BYTETracker
from spacemit_cv import AGVDetection,BYTETracker

def main(args):
   
   tracker = BYTETracker(frame_rate=30)
   frame_id = 0
   results = []

   detector = AGVDetection(args.model)
   detector.warm_up()

   cap = cv2.VideoCapture(1)

   while True:
       ret, frame = cap.read()

       if ret:
            outputs = detector.infer_track(frame)
            online_targets = tracker.update(outputs,[frame.shape[0],frame.shape[1]],[frame.shape[0],frame.shape[1]])
    

            online_tlwhs = []
            online_ids = []
            online_scores = []
            for t in online_targets:
                tlwh = t.tlwh
                tid = t.track_id
                vertical = tlwh[2] / tlwh[3] > 1.6
                if tlwh[2] * tlwh[3] > args.min_box_area and not vertical:
                    online_tlwhs.append(tlwh)
                    online_ids.append(tid)
                    online_scores.append(t.score)
                    cv2.rectangle(frame, (int(tlwh[0]), int(tlwh[1])), (int(tlwh[0] + tlwh[2]), int(tlwh[1] + tlwh[3])), (0, 255, 0), 2)
                    cv2.putText(frame, str(tid),(int(tlwh[0]), int(tlwh[1])),0, 5e-3 * 200, (0, 255, 0),2)
   
            cv2.imshow("Tracking", frame)
            ch = cv2.waitKey(1)
            if ch == 27 or ch == ord("q") or ch == ord("Q"):
                break
       else:
            break
       frame_id += 1






if __name__ == '__main__':

    # 解析命令行参数
    parser = argparse.ArgumentParser(description='YOLOv8 ONNX Inference')
    # detection args
    parser.add_argument('--model', type=str, default='spacemit_cv/model/yolov8n.q.onnx', help='Path to the YOLOv8 ONNX model')
    # tracking args
    parser.add_argument("--track_thresh", type=float, default=0.5, help="tracking confidence threshold")
    parser.add_argument("--track_buffer", type=int, default=30, help="the frames for keep lost tracks")
    parser.add_argument("--match_thresh", type=float, default=0.8, help="matching threshold for tracking")
    parser.add_argument('--min-box-area', type=float, default=10, help='filter out tiny boxes')    
    parser.add_argument("--mot20", dest="mot20", default=False, action="store_true", help="test mot20.")
     
    args = parser.parse_args()

    main(args)